package com.koicarp.agent.provider;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;

import com.koicarp.agent.common.KoiException;
import com.koicarp.agent.common.constant.CommonConstant;
import com.koicarp.agent.provider.constant.ProviderErrMsgConstant;
import org.springframework.util.CollectionUtils;


import cn.hutool.core.util.StrUtil;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageType;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import lombok.extern.slf4j.Slf4j;

/**
 * 用户聊天的缓存
 * @author shen
 * @date    2025年3月6日
 * 下午3:12:33
 */
@Slf4j
public class ChatMemeoryCache implements ChatMemoryStore {

	/**
	 *会话id，消息集合
	 */
	 private static final Map<Object, List<ChatMessage>> store = new HashMap<>();
    
    /**
     * 应用id，用户id，会话id
     */
    private static final Map<Integer, Map<Integer, Integer>> appUserMagMap = new HashMap<>();
    
    /**
     * 判断提示词有没有初始化
     */
    private static final Map<Object, Boolean> initSystemMessageStore = new HashMap<>();

    public static void clean(Object memoryId) {
        log.info("clean message memory store to: {}", memoryId);
        store.remove(memoryId);
    }

    /**
     * @Description: 初始化系统缓存
     * @param memoryId
     * @param message void
     * @date 2025-03-11 
     * @author shen
     */
    public static List<ChatMessage> initSystem(Object memoryId, SystemMessage message) {
        Boolean isInitSystemMessage = initSystemMessageStore.get(memoryId);
        if (isInitSystemMessage != null && isInitSystemMessage) {
            return store.get(memoryId);
        }

        initSystemMessageStore.put(memoryId, true);
        return new ArrayList<>(Arrays.asList(message));     
    }

    /**
     * @Description: 获取用户在这个应用下的会话id
     * @param appId
     * @param userId
     * @return Integer
     * @date 2025-03-06 
     * @author shen
     */
    public static Integer getConversationId(Integer appId, Integer userId) {
    	Map<Integer, Integer> map = appUserMagMap.get(appId);
    	
    	if(CollectionUtils.isEmpty(map) || Objects.isNull(map)) {
    		return null;
    	}
    	Integer cacheConversationId = map.get(userId);   	
        return cacheConversationId;
    }
    
    /**
     * @Description: 切换用户的对话缓存
     * @param appId
     * @param userId
     * @param newConversationId
     * @date 2025-03-06 
     * @author shen
     */
    public static void init(Integer appId, Integer userId, Integer newConversationId, List<ChatMessage> msgs, String promt) {    
    	//如果没有切换对话，缓存里本身就有，不用进行处理
    	Integer oldConversationId = ChatMemeoryCache.getConversationId(appId, userId);
    	if(newConversationId.equals(oldConversationId)) return;
    	
    	//此处是为了保证每次systemMsg都在送往大模型传参的第一个
    	List<ChatMessage> messages = new ArrayList<>();
    	//处理提示词
    	if (StrUtil.isNotBlank(promt)) {
            // initialize chat memory
            SystemMessage message = new SystemMessage(promt);
            List<ChatMessage> initSystem = new ArrayList<>(Arrays.asList(message));
            messages.addAll(initSystem);
        }else {
        	log.error("聊天前更新缓存时appId[{}]没有提示词，抛出异常终止聊天过程", appId);
			throw new KoiException(ProviderErrMsgConstant.CHAT_APP_PROMPT_ERR);
		}
    	messages.addAll(msgs);
    	
    	
    	//1-开启新对话
    	if(Objects.isNull(oldConversationId)) {
        	//更新应用下用户和会话的缓存
    		Map<Integer, Integer> map = appUserMagMap.get(appId);
    		if(CollectionUtils.isEmpty(map) || Objects.isNull(map)) {
    			map = new HashMap<>();
    		}
        	map.put(userId, newConversationId);
        	appUserMagMap.put(appId, map);
    		store.put(newConversationId, messages);   	
    	}else {//2-切换对话
    		if(!newConversationId.equals(oldConversationId)) {
    	    	//更新应用下用户和会话的缓存
    			Map<Integer, Integer> map = appUserMagMap.get(appId);
    	    	map.put(userId, newConversationId);
    			//删除老的，添加新的 	
    	    	store.remove(oldConversationId); //更新会话与消息的缓存
    	    	store.put(newConversationId, messages); 
    	    	log.info("用户从老对话[{}]，切换到新对话id[{}],缓存初始化新的聊天成功",oldConversationId, newConversationId);
    		}
    		
    	}  	  	
    }
    
    /**
     * @Description: 更新缓存
     * @param memoryId
     * @param messages void
     * @date 2025-03-06 
     * @author shen
     */
    public static void update(Object memoryId, List<ChatMessage> messages) {
        List<ChatMessage> list = store.get(memoryId);
        if (CollectionUtils.isEmpty(list)) {
            store.put(memoryId, messages);
        } else {
            list.addAll(messages);
        }
    }
    
    public static Map<Integer, Map<Integer, Integer>> getAllUser() {
       	return appUserMagMap;
    }
    
    public static Map<Object, List<ChatMessage>> getAllMsg() {
       	return store;
    }
    
    //----------------继承的实现方法----------------------------
    
    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        List<ChatMessage> list = store.get(memoryId);
        if (list == null) {
            return new ArrayList<>();
        }
        return list;
    }

	@Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
    	ChatMessage chatMessage = messages.get(messages.size()- 1);
    	String text = chatMessage.text();
    	ChatMessageType type = chatMessage.type();
    	Class<? extends ChatMessage> messageClass = type.messageClass();
    	//20250313 shen 去掉AI回答的think标签
    	if(messageClass.equals(AiMessage.class)) {
    		messages.remove(chatMessage);
    		String subString = this.getSubString(CommonConstant.THINK_BEFORE, CommonConstant.THINK_AFTER, text);
        	text = subString;
        	messages.add(new AiMessage(text));
    	}
    	
    	store.put(memoryId, messages);
    }
	
	private  String getSubString(String beginStr, String endStr, String source) {		 		 		
		if(!source.contains(beginStr)) return source;
		
        int strStartIndex = source.indexOf(beginStr);
        int strEndIndex = source.indexOf(endStr) + endStr.length();
        String subStrToRemove = source.substring(strStartIndex, strEndIndex);        
        String replace = source.replace(subStrToRemove, "");        
        return replace;
	}

    @Override
    public void deleteMessages(Object memoryId) {
        store.remove(memoryId);
    }
}